import time

import numpy as np
import gurobipy as gp
from gurobipy import GRB

import verify_piecewise_utils
from NN import cdinn2
from prob import *

# case_name = "C11"

def block_space(dim_min, dim_max):
    n = len(dim_min)

    sub_blocks = []

    # 递归地对当前维度进行划分
    def split_block(current_min, current_max, index):
        if index == n:
            sub_blocks.append([current_min, current_max])
            return
        mid = (current_min[index] + current_max[index]) / 2

        # 左子块
        left_max = current_max[:]
        left_max[index] = mid
        split_block(current_min, left_max, index + 1)

        # 右子块
        right_min = current_min[:]
        right_min[index] = mid
        split_block(right_min, current_max, index + 1)

    # 开始分割
    split_block(dim_min, dim_max, 0)
    return sub_blocks

def barrier_f_without_der(m, x0, wb_list, wb_length, useVars):
    x_layer = x0
    for i in range(wb_length):
        W = wb_list[i * 2]
        b = wb_list[i * 2 + 1]
        if useVars:
            y_layer = m.addMVar((W.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
            E = np.identity(y_layer.shape[0])
            expr = W @ x_layer + b - E @ y_layer
            m.addConstr(expr == 0)
        else:
            y_layer = W @ x_layer + b

        if i != wb_length - 1:
            if useVars:
                z_layer = m.addMVar((y_layer.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
                # z = activate(y)
                for j in range(y_layer.shape[0]):
                    m.addConstr(z_layer[j] == gp.max_(y_layer[j], 0.0))
            else:
                z_layer = y_layer
                for j in range(y_layer.shape[0]):
                    z_layer[j] = max(y_layer[j], 0.0)
            x_layer = z_layer
        else:
            x_layer = y_layer

        if useVars:
            m.update()
    return x_layer


def barrier_Bf_nonlinear(wb_list, wb_length, min_bound, max_bound):
    m = gp.Model()
    m.setParam('Outputflag', 0)
    m.setParam('NonConvex', 2)
    A = m.addVars(DIM, name="A")
    b_l = m.addVar(name="b_l")
    b = m.addVar(name="b")
    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
    x0[0].lb = min_bound[0]
    x0[0].ub = max_bound[0]
    x0[1].lb = min_bound[1]
    x0[1].ub = max_bound[1]
    x0[2].lb = min_bound[2]
    x0[2].ub = max_bound[2]
    x0[3].lb = min_bound[3]
    x0[3].ub = max_bound[3]
    x1 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x1', lb=-np.inf, ub=np.inf)  # x1 = F(kx+b; x0)
    x1[0].lb = domain_min[0]
    x1[0].ub = domain_max[0]
    x1[1].lb = domain_min[1]
    x1[1].ub = domain_max[1]
    x1[2].lb = domain_min[2]
    x1[2].ub = domain_max[2]
    x1[3].lb = domain_min[3]
    x1[3].ub = domain_max[3]

    xx1 = m.addVar()
    xx3 = m.addVar()
    m.addConstr(xx1 == x0[1] * x0[1])
    m.addConstr(xx3 == x0[3] * x0[3])
    m.addConstr(x1[0] == x0[0] + time_step * (- x0[0] + (xx1 * x0[1]) - 3 * (x0[2] * x0[3])))
    m.addConstr(x1[1] == x0[1] + time_step * (- x0[0] - (xx1 * x0[1])))
    m.addConstr(x1[2] == x0[2] + time_step * ((x0[0] * x0[3]) - x0[2]))
    m.addConstr(x1[3] == x0[3] + time_step * ((x0[0] * x0[2]) - (xx3 * x0[3])))

    Bf = barrier_f_without_der(m, x1, wb_list, wb_length, True)
    m.addConstr(A[0] * x1[0] + A[1] * x1[1] + A[2] * x1[2] + A[3] * x1[3] + b_l <= Bf)
    m.addConstr(A[0] * x1[0] + A[1] * x1[1] + A[2] * x1[2] + A[3] * x1[3] + b_l + b >= Bf)
    m.setObjective(b, GRB.MINIMIZE)
    m.optimize()
    if m.status == 2:
        return True, A[0].x, A[1].x, A[2].x, A[3].x, b_l.x + b.x
    else:
        return False, None, None, None, None, None

def barrier_f_linear(wb_list, wb_length, min_bound, max_bound):
    m = gp.Model()
    m.setParam('Outputflag', 0)
    m.setParam('NonConvex', 2)
    A = m.addVars(DIM, name="A")
    b_l = m.addVar(name="b_l")
    b = m.addVar(name="b")
    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
    x0[0].lb = min_bound[0]
    x0[0].ub = max_bound[0]
    x0[1].lb = min_bound[1]
    x0[1].ub = max_bound[1]
    x0[2].lb = min_bound[2]
    x0[2].ub = max_bound[2]
    x0[3].lb = min_bound[3]
    x0[3].ub = max_bound[3]
    B = barrier_f_without_der(m, x0, wb_list, wb_length, True)
    m.addConstr(A[0] * x0[0] + A[1] * x0[1] + A[2] * x0[2] + A[3] * x0[3] + b_l <= B)
    m.addConstr(A[0] * x0[0] + A[1] * x0[1] + A[2] * x0[2] + A[3] * x0[3] + b_l + b >= B)
    m.setObjective(b, GRB.MINIMIZE)
    m.optimize()
    if m.status == 2:
        return True, A[0].x, A[1].x, A[2].x, A[3].x, b.x
    else:
        return False, None, None, None, None, None

# min hNN-gNN
def relax_verify_initCond(wb_1_list, wb_2_list, wb_length, model):
    cex_init = []
    res_init = []

    init_len = 0.5
    cnt = 0
    cex_idx_list = []

    while True:
        relax_flag = False

        new_cex_list = []
        if len(cex_idx_list) == 0:
            cex_idx_list.append([init_min, init_max])

        # print("init_cex_idx_list_num:", len(cex_idx_list))

        for cex_idx in cex_idx_list:
            block_min, block_max = cex_idx
            x0_l = block_min[0]
            x1_l = block_min[1]
            x2_l = block_min[2]
            x3_l = block_min[3]
            x0_r = block_max[0]
            x1_r = block_max[1]
            x2_r = block_max[2]
            x3_r = block_max[3]

            m = gp.Model()
            m.setParam('Outputflag', 0)
            x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)

            m.addConstr(x0[0] >= x0_l)
            m.addConstr(x0[0] <= x0_r)
            m.addConstr(x0[1] >= x1_l)
            m.addConstr(x0[1] <= x1_r)
            m.addConstr(x0[2] >= x2_l)
            m.addConstr(x0[2] <= x2_r)
            m.addConstr(x0[3] >= x3_l)
            m.addConstr(x0[3] <= x3_r)

            hNN = barrier_f_without_der(m, x0, wb_2_list, wb_length, True)

            # hat_gNN = Ax + b = hNN upper bound
            feasible, A1, A2, A3, A4, b = barrier_f_linear(wb_1_list, wb_length, [x0_l, x1_l, x2_l, x3_l], [x0_r, x1_r, x2_r, x3_r])
            if not feasible:
                print("unfeasible init barrier_f_linear")
                continue

            m.setObjective(hNN - (A1 * x0[0] + A2 * x0[1] + A3 * x0[2] + A4 * x0[3] + b), GRB.MINIMIZE)
            m.optimize()

            if m.status == 2:
                if m.objVal < 0:
                    if model(torch.tensor(np.array(x0.x), dtype=torch.float64)) > 0:
                        cex_init.append(np.array(x0.x))
                        res_init.append(m.objVal)
                        break
                    else:
                        relax_flag = True
                        new_cex_list += block_space(block_min, block_max)

        if relax_flag and len(cex_init) == 0 and cnt<cnt_num:
            init_len /= 2
            cnt += 1
            print("init_len:", init_len)
            cex_idx_list = new_cex_list

        else:
            break

    return cex_init, res_init

# min gNN-hNN
def relax_verify_unsafeCond(wb_1_list, wb_2_list, wb_length, model):
    cex_unsafe = []
    res_unsafe = []

    unsafe_len = 1
    cnt = 0
    cex_idx_list = []
    while True:
        relax_flag = False

        new_cex_list = []
        if len(cex_idx_list) == 0:
            cex_idx_list.append([unsafe_min, unsafe_max])

        # print("unsafe_cex_idx_list_num:", len(cex_idx_list))

        for cex_idx in cex_idx_list:
            block_min, block_max = cex_idx
            x0_l = block_min[0]
            x1_l = block_min[1]
            x2_l = block_min[2]
            x3_l = block_min[3]
            x0_r = block_max[0]
            x1_r = block_max[1]
            x2_r = block_max[2]
            x3_r = block_max[3]

            m = gp.Model()
            m.setParam('Outputflag', 0)
            x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)

            m.addConstr(x0[0] >= x0_l)
            m.addConstr(x0[0] <= x0_r)
            m.addConstr(x0[1] >= x1_l)
            m.addConstr(x0[1] <= x1_r)
            m.addConstr(x0[2] >= x2_l)
            m.addConstr(x0[2] <= x2_r)
            m.addConstr(x0[3] >= x3_l)
            m.addConstr(x0[3] <= x3_r)

            if unsafe_shape == 2:
                p0 = (unsafe_min[0] + unsafe_max[0]) / 2.0
                p1 = (unsafe_min[1] + unsafe_max[1]) / 2.0
                p2 = (unsafe_min[2] + unsafe_max[2]) / 2.0
                p3 = (unsafe_min[3] + unsafe_max[3]) / 2.0
                r = p0 - unsafe_min[0]
                m.addConstr((x0[0] - p0) * (x0[0] - p0) + (x0[1] - p1) * (x0[1] - p1)
                            + (x0[2] - p2) * (x0[2] - p2) + (x0[3] - p3) * (x0[3] - p3) <= r * r)

            gNN = barrier_f_without_der(m, x0, wb_1_list, wb_length, True)

            # hat_hNN = Ax + b = hNN upper bound
            feasible, A1, A2, A3, A4, b = barrier_f_linear(wb_2_list, wb_length, [x0_l, x1_l, x2_l, x3_l], [x0_r, x1_r, x2_r, x3_r])
            if not feasible:
                print("unfeasible unsafe barrier_f_linear")
                continue
            m.setObjective(gNN - (A1 * x0[0] + A2 * x0[1] + A3 * x0[2] + A4 * x0[3] + b), GRB.MINIMIZE)
            m.optimize()

            # print("minimize unsafe value:", m.objVal)
            if m.status == 2:
                if m.objVal < 0:
                    if model(torch.tensor(np.array(x0.x), dtype=torch.float64)) < 0:
                        cex_unsafe.append(np.array(x0.x))
                        res_unsafe.append(m.objVal)
                        break
                    else:
                        relax_flag = True
                        new_cex_list += block_space(block_min, block_max)

        if relax_flag and len(cex_unsafe)==0 and cnt<cnt_num:
            unsafe_len /= 2
            cnt += 1
            print("unsafe_len:", unsafe_len)
            cex_idx_list = new_cex_list

        else:
            break
    return cex_unsafe, res_unsafe

# max ( (gNN(f)-hNN(f)) - (1-lambda) * (gNN-hNN) )
def global_verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len, model):
    cex_domain = []
    res_domain = []

    domain_len = 4
    cnt = 0

    x0_min = domain_min[0]
    x0_max = domain_max[0]
    x1_min = domain_min[1]
    x1_max = domain_max[1]
    x2_min = domain_min[2]
    x2_max = domain_max[2]
    x3_min = domain_min[3]
    x3_max = domain_max[3]

    cex_idx_list = []
    while True:
        relax_flag = False

        new_cex_list = []
        if len(cex_idx_list) == 0:
            cex_idx_list.append([domain_min, domain_max])

        print("domain_cex_idx_list_num:", len(cex_idx_list))

        for cex_idx in cex_idx_list:
            block_min, block_max = cex_idx
            x0_l = block_min[0]
            x1_l = block_min[1]
            x2_l = block_min[2]
            x3_l = block_min[3]
            x0_r = block_max[0]
            x1_r = block_max[1]
            x2_r = block_max[2]
            x3_r = block_max[3]

            m = gp.Model()
            m.setParam('Outputflag', 0)
            m.setParam('NonConvex', 2)
            # m.setParam('InfUnbdInfo', 1)

            x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
            x0[0].lb = x0_l
            x0[0].ub = x0_r
            x0[1].lb = x1_l
            x0[1].ub = x1_r
            x0[2].lb = x2_l
            x0[2].ub = x2_r
            x0[3].lb = x3_l
            x0[3].ub = x3_r

            x1 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x1', lb=-np.inf, ub=np.inf)  # x1 = F(kx+b; x0)
            x1[0].lb = x0_min
            x1[0].ub = x0_max
            x1[1].lb = x1_min
            x1[1].ub = x1_max
            x1[2].lb = x2_min
            x1[2].ub = x2_max
            x1[3].lb = x3_min
            x1[3].ub = x3_max

            xx1 = m.addVar()
            xx3 = m.addVar()
            m.addConstr(xx1 == x0[1] * x0[1])
            m.addConstr(xx3 == x0[3] * x0[3])
            m.addConstr(x1[0] == x0[0] + time_step * (- x0[0] + (xx1 * x0[1]) - 3 * (x0[2] * x0[3])))
            m.addConstr(x1[1] == x0[1] + time_step * (- x0[0] - (xx1 * x0[1])))
            m.addConstr(x1[2] == x0[2] + time_step * ((x0[0] * x0[3]) - x0[2]))
            m.addConstr(x1[3] == x0[3] + time_step * ((x0[0] * x0[2]) - (xx3 * x0[3])))

            gNN = barrier_f_without_der(m, x0, wb_1_list, wb_length, True)
            hNN_f = barrier_f_without_der(m, x1, wb_2_list, wb_length, True)

            feasible, gNN_f_A0, gNN_f_A1, gNN_f_A2, gNN_f_A3, gNN_b = barrier_f_linear(wb_1_list, wb_length, [x0_l, x1_l, x2_l, x3_l], [x0_r, x1_r, x2_r, x3_r])
            if not feasible:
                continue
            feasible, hNN_A0, hNN_A1, hNN_A2, hNN_A3, hNN_b = barrier_f_linear(wb_2_list, wb_length, [x0_l, x1_l, x2_l, x3_l], [x0_r, x1_r, x2_r, x3_r])
            if not feasible:
                continue
            m.setObjective( (gNN_f_A0*x1[0] + gNN_f_A1*x1[1] + gNN_f_A2*x1[2] + gNN_f_A3*x1[3] + gNN_b)
                            -hNN_f
                            - (1-alpha) * gNN
                            + (1-alpha) * (hNN_A0*x0[0] + hNN_A1*x0[1] + hNN_A2*x0[2] + hNN_A3*x0[3] + hNN_b),
                            GRB.MAXIMIZE)

            m.optimize()
            if m.status == 2:
                if m.objVal > 0:
                    # print(f"m.objVal:{m.objVal}, model(x0.x):{model(vector_field(torch.tensor([np.array(x0.x)]))) - (1-alpha) * model(torch.tensor(np.array(x0.x)))}")
                    if model(vector_field(torch.tensor([np.array(x0.x)]))) - (1 - alpha) * model(
                            torch.tensor(np.array(x0.x))) > 0:
                        cex_domain.append(np.array(x0.x))
                        res_domain.append(m.objVal)
                        # print(np.array(x0.x))
                        # print(cex_domain)
                        # print(m.objVal)
                        # print(res_domain)
                        break
                    else:
                        relax_flag = True
                        new_cex_list += block_space(block_min, block_max)

        if relax_flag and len(cex_domain) == 0 and cnt < cnt_num:
            domain_len /= 2
            cnt += 1
            print("domain_len:", domain_len)
            cex_idx_list = new_cex_list
        else:
            break

    return cex_domain, res_domain


def verif(model, x_len):
    wb_1_list = []
    wb_2_list = []
    params = model.state_dict()
    i = 0
    for k in params:
        # print(k)
        # 'hidden_layers1.0.weight' 'hidden_layers1.0.bias' 'hidden_layers2.0.weight' 'hidden_layers2.0.bias'
        if 'layers' in k:
            continue
        if i % 4 < 2:
            wb_1_list.append(params[k].detach().numpy())
        else:
            wb_2_list.append(params[k].detach().numpy())
        i += 1

    wb_length = int(len(wb_1_list) / 2)
    # time_init_start = time.time()
    cex_init, res_init = relax_verify_initCond(wb_1_list, wb_2_list, wb_length, model)
    # time_init_end = time.time()
    # print("init cond time:", time_init_end - time_init_start)

    # time_unsafe_start = time.time()
    cex_unsafe, res_unsafe = relax_verify_unsafeCond(wb_1_list, wb_2_list, wb_length, model)
    # time_unsafe_end = time.time()
    # print("unsafe cond time:", time_unsafe_end - time_unsafe_start)

    # time_domain_start = time.time()
    cex_domain, res_domain = global_verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len, model)
    # time_domain_end = time.time()
    # print("domain cond time:", time_domain_end - time_domain_start)

    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain

